from openai_api import ChatCompletionClient
import os
from dotenv import load_dotenv
from paddleocr import PaddleOCR
import json
import markdown
import sys

load_dotenv("config.env")

def read_json_file(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        return data
    except FileNotFoundError:
        print("错误: 文件未找到，请检查路径文件路径。")
    except json.JSONDecodeError:
        print("错误: JSON格式不正确。")

def read_prompt(prompt_path):
    article_correction_system_prompt_path = os.path.join(prompt_path, "article_correction_system_prompt.txt")
    article_correction_user_prompt_path = os.path.join(prompt_path, "article_correction_user_prompt.txt")

    with open(article_correction_system_prompt_path, "r", encoding="utf-8") as f:
        article_correction_system = f.read()

    with open(article_correction_user_prompt_path, "r", encoding="utf-8") as f:
        article_correction_user = f.read()

    return article_correction_system, article_correction_user

if __name__ == '__main__':
    prompt_path = "./prompt"
    article_correction_system, article_correction_user = read_prompt(prompt_path)
    
    article_client = ChatCompletionClient(
        base_url=os.getenv("base_url"),
        api_key=os.getenv("api_key"),
        model_name="deepseek-ai/DeepSeek-V3",
        system=article_correction_system
    )

    while True:
        if len(sys.argv) != 2:
            print("Usage: python main.py <param1>")
            sys.exit(1)
        image_name=sys.argv[1]
        image_path=image_name
        print(image_name)
        ocr = PaddleOCR(
            use_doc_orientation_classify=True,
            use_doc_unwarping=True,
            use_textline_orientation=True
        )
        result = ocr.predict(input=image_path)
        for res in result:
            res.print()
            res.save_to_img("cache")
            res.save_to_json("cache")
        k=len(image_name)
        json_name=image_name[7:k-4]
        json_name="cache\\"+json_name+'_res.json'
        json_dict=read_json_file(json_name)
        text_list=json_dict.get("rec_texts")
        article_title=text_list[0]
        k=len(text_list)
        text=""
        for i in range(1,k-1):
            text+=text_list[i]
            text+="\n"
        # 计算文章长度
        article_length = len(text)-(k-2)
        # 替换模板中的变量
        input_text = article_correction_user.format(
            article_title=article_title,
            article_length=article_length,
            text=text
        )
        out_put = article_client.create_chat_completion(prompt=input_text,stream=True)
        name1=markdown.markdown("学生姓名/代号："+image_name[7:k-4])
        output1=markdown.markdown(out_put)
        m=input("回车导出结果......")
        f_data=open("result.html","a+",encoding="utf-8")
        f_data.write(name1+"\n")
        f_data.write(output1+"\n\n")
        f_data.close()
        break

